// Copyright 2020 Makani Technologies LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Convert a set of .pcap files recording AIO network traffic to an
// HDF5 file with the structure:
//
//   /messages/kAioNode<source>/kMessageType<type>
//
// Each dataset consists of the following fields:
//
//   capture: Information from pcap (source, destination, timestamp).
//   header:  AioHeader for that packet.
//   message: Payload.
//
// The payload formatting is copied from an HDF5 file generated by
// generate_hdf5_format.py and compiled into the pcap_to_hdf5 binary.
//
// Usage:
//
//   ./lib/pcap_to_hdf5/pcap_to_hdf5 -output_file=output.h5 1.pcap 2.pcap ...
//
// TODO: Update pcap_to_hdf5 so that it does not need to be
// built from the same commit as the pcap file (i.e. by putting more
// information in the formatting file (see MessageLog) and taking the
// format file as a command line option).

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <netinet/if_ether.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <pcap/pcap.h>
#include <pcap/sll.h>
#include <stdint.h>

#include <limits>

#include "avionics/common/aio_header.h"
#include "avionics/common/aio_version_code.h"
#include "avionics/common/network_config.h"
#include "avionics/common/pack_aio_header.h"
#include "avionics/network/aio_node.h"
#include "avionics/network/message_type.h"
#include "control/control_params.h"
#include "control/pack_control_telemetry.h"
#include "control/system_params.h"
#include "lib/json_load/load_params.h"
#include "lib/pcap_reader/pcap_reader.h"
#include "lib/pcap_to_hdf5/capture_info.h"
#include "lib/pcap_to_hdf5/message_log.h"
#include "sim/pack_sim_telemetry.h"
#include "sim/sim_params.h"

DEFINE_string(output_file, "", "The path of the output log file.");
DEFINE_bool(write_empty_log, false, "Whether to write an empty log.");
DEFINE_bool(ignore_version, false, "Ignore AIO version check.");
DEFINE_bool(dedup, true, "Deduplicate AIO messages.");

namespace {

// Checks that the output file is non-empty.
bool ValidateOutputFile(const char * /* flag_name */,
                        const std::string &value) {
  return value.size() > 0;
}
bool dummy __attribute__((unused)) =
    google::RegisterFlagValidator(&FLAGS_output_file, &ValidateOutputFile);

// Writes a specific subset of the possible messages for use in unit tests.
void WriteEmptyLog() {
  const struct {
    AioNode source;
    MessageType type;
    size_t size;
  } messages[] = {
      {kAioNodeControllerA, kMessageTypeControlTelemetry,
       PACK_CONTROLTELEMETRY_SIZE},
      {kAioNodeControllerA, kMessageTypeControlDebug,
       PACK_CONTROLTELEMETRY_SIZE},
      {kAioNodeSimulator, kMessageTypeSimTelemetry, PACK_SIMTELEMETRY_SIZE}};

  lib::pcap_to_hdf5::MessageLog log(FLAGS_output_file);

  for (int32_t i = 0; i < ARRAYSIZE(messages); ++i) {
    std::vector<uint8_t> buffer(messages[i].size, 0U);
    log.WriteWithSpoofedHeaders(messages[i].source, messages[i].type, 0U,
                                buffer.data(),
                                static_cast<int32_t>(buffer.size()));
  }

  lib::pcap_reader::PcapReader pcap_reader;
  log.WriteLogInfo(pcap_reader);
}

bool HandleIp(lib::pcap_reader::PcapReader *reader, BadPacketInfo *bad,
              lib::pcap_to_hdf5::MessageLog *log) {
  switch (reader->ReadIp()) {
    case lib::pcap_reader::IpStatus::kBadEthernetHeader:
      bad->reason = kBadPacketReasonEthernetHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kBadEthernetProtocol:
      bad->reason = kBadPacketReasonEthernetProtocol;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kBadIpHeader:
      bad->reason = kBadPacketReasonIpHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kBadLinkProtocol:
      bad->reason = kBadPacketReasonLinkProtocol;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kBadPcapHeader:
      bad->reason = kBadPacketReasonPcapHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kBadSllHeader:
      bad->reason = kBadPacketReasonSllHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::IpStatus::kFragmented:
      return false;
    case lib::pcap_reader::IpStatus::kValid:
      return true;
    default:
      return false;
  }
}

bool HandleUdp(lib::pcap_reader::PcapReader *reader, BadPacketInfo *bad,
               lib::pcap_to_hdf5::MessageLog *log) {
  switch (reader->ReadUdp()) {
    case lib::pcap_reader::UdpStatus::kBadIpProtocol:
      bad->reason = kBadPacketReasonIpProtocol;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::UdpStatus::kBadUdpHeader:
      bad->reason = kBadPacketReasonUdpHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::UdpStatus::kValid:
      return true;
    default:
      return false;
  }
}

bool HandleAio(lib::pcap_reader::PcapReader *reader, BadPacketInfo *bad,
               lib::pcap_to_hdf5::MessageLog *log) {
  switch (reader->ReadAio(FLAGS_ignore_version)) {
    case lib::pcap_reader::AioStatus::kBadAioHeader:
      bad->reason = kBadPacketReasonAioHeader;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kBadAioVersion:
      assert(!FLAGS_ignore_version);
      bad->reason = kBadPacketReasonAioVersion;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kBadIpProtocol:
      bad->reason = kBadPacketReasonIpProtocol;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kBadMessageType:
      bad->reason = kBadPacketReasonAioType;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kBadSource:
      bad->reason = kBadPacketReasonAioSource;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kBadUdpPort:
      bad->reason = kBadPacketReasonAioPort;
      log->WriteBadPacketInfo(*bad);
      return false;
    case lib::pcap_reader::AioStatus::kDuplicate:
      return !FLAGS_dedup;
    case lib::pcap_reader::AioStatus::kValid:
      return true;
    default:
      return false;
  }
}

bool HandlePcapMessage(lib::pcap_reader::PcapReader *reader,
                       lib::pcap_to_hdf5::MessageLog *log) {
  // Build bad packet info.
  BadPacketInfo bad;
  memset(&bad, 0, sizeof(bad));
  bad.capture_header.tv_sec = reader->GetHeader().ts.tv_sec;
  bad.capture_header.tv_usec = reader->GetHeader().ts.tv_usec;

  // Handle IP layer.
  if (!HandleIp(reader, &bad, log)) {
    return true;
  }

  // Update capture header with IP information.
  bad.capture_header.source_address = reader->GetIp().GetSource();
  bad.capture_header.destination_address = reader->GetIp().GetDestination();

  // Handle UDP layer.
  if (!HandleUdp(reader, &bad, log)) {
    return true;
  }
  if (reader->GetUdp().GetDestination() != UDP_PORT_AIO) {
    // Proceed to next packet if this is not an AIO packet.
    return true;
  }

  // Handle AIO layer.
  if (!HandleAio(reader, &bad, log)) {
    return true;
  }

  // Accumate timing stats.
  log->CalcTimingStats(reader->GetAio().GetHeader(), bad.capture_header,
                       reader->GetAio().GetPayload());

  // Write valid AIO packet.
  log->WriteData(bad.capture_header, reader->GetAio().GetHeader(),
                 reader->GetUdp().GetPayload(),
                 reader->GetUdp().GetPayloadLength());
  return true;
}

}  // namespace

#if !defined(MAKANI_TEST)

int main(int argc, char *argv[]) {
  google::InitGoogleLogging(argv[0]);
  google::SetUsageMessage(
      "\nUsage: pcap_to_hdf5 -output_file=output.h5 1.pcap [2.pcap] ...\n"
      "       pcap_to_hdf5 -output_file=output.h5 -write_empty_log");
  google::ParseCommandLineFlags(&argc, &argv, true);

  json_load::LoadSystemParams(GetSystemParamsUnsafe());
  json_load::LoadControlParams(GetControlParamsUnsafe());
  json_load::LoadSimParams(GetSimParamsUnsafe());

  if (FLAGS_write_empty_log) {
    if (argc > 1) {
      google::ShowUsageWithFlagsRestrict(argv[0], "pcap_to_hdf5");
      return 1;
    }
    WriteEmptyLog();
    return 0;
  }

  if (argc < 2) {
    google::ShowUsageWithFlagsRestrict(argv[0], "pcap_to_hdf5");
    return 1;
  }

  lib::pcap_reader::PcapReader pcap_reader;
  lib::pcap_to_hdf5::MessageLog log(FLAGS_output_file);

  // TODO: Add basic statistics to the log file for packets
  // that are not AIO traffic.
  if (pcap_reader.ForEach(argc - 1, argv + 1, HandlePcapMessage, &log) < 0) {
    LOG(FATAL) << "Pcap error: " << pcap_reader.GetError();
  }
  log.WriteLogInfo(pcap_reader);
}

#endif  // !defined (MAKANI_TEST)
